package com.joker.demo;

import javax.sound.sampled.*;
import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.File;
import java.io.IOException;
import java.util.concurrent.CountDownLatch;

public class MicrophoneRecorder {
    private TargetDataLine line;
    private ByteArrayOutputStream out;
    private boolean isRecording = false;
    private AudioFormat format;

    public static void main(String[] args) throws Exception {
        // 使用示例
        MicrophoneRecorder recorder = new MicrophoneRecorder();

        // 录制5秒并保存为WAV文件
        byte[] audioBytes = recorder.record(5);

        // 保存为WAV文件
        recorder.saveAsWav(audioBytes, "recorded_audio1.wav");
        System.out.println("录音已保存为 recorded_audio.wav");

//        // 也可以直接录制并保存
//        recorder.recordAndSave(5, "direct_record.wav");
//        System.out.println("直接录音已保存为 direct_record.wav");
    }

    public byte[] record(int durationSeconds) throws Exception {
        // 设置音频格式：16kHz采样率，16位样本，单声道，有符号，小端序
        format = new AudioFormat(16000, 16, 1, true, false);
        DataLine.Info info = new DataLine.Info(TargetDataLine.class, format);

        if (!AudioSystem.isLineSupported(info)) {
            throw new Exception("音频格式不支持");
        }

        line = (TargetDataLine) AudioSystem.getLine(info);
        line.open(format);
        line.start();

        out = new ByteArrayOutputStream();
        isRecording = true;
        CountDownLatch latch = new CountDownLatch(1);

        // 启动录音线程
        new Thread(() -> {
            byte[] buffer = new byte[1024];
            while (isRecording) {
                int bytesRead = line.read(buffer, 0, buffer.length);
                if (bytesRead > 0) {
                    out.write(buffer, 0, bytesRead);
                }
            }
            latch.countDown();
        }).start();

        // 录音指定时长
        Thread.sleep(durationSeconds * 1000);
        stopRecording();
        latch.await();

        return out.toByteArray();
    }

    /**
     * 将录音数据保存为WAV文件
     * @param audioData 原始音频数据
     * @param filename 文件名
     * @throws IOException 文件操作异常
     */
    public void saveAsWav(byte[] audioData, String filename) throws IOException {
        if (format == null) {
            throw new IllegalStateException("请先调用record方法进行录音");
        }

        // 创建音频输入流
        ByteArrayInputStream bais = new ByteArrayInputStream(audioData);
        AudioInputStream audioInputStream = new AudioInputStream(bais, format, audioData.length / format.getFrameSize());

        // 保存为WAV文件
        AudioSystem.write(audioInputStream, AudioFileFormat.Type.WAVE, new File(filename));
        audioInputStream.close();
    }

    /**
     * 直接录制并保存为WAV文件
     * @param durationSeconds 录音时长（秒）
     * @param filename 文件名
     * @throws Exception 录音或文件操作异常
     */
    public void recordAndSave(int durationSeconds, String filename) throws Exception {
        byte[] audioData = record(durationSeconds);
        saveAsWav(audioData, filename);
    }

    /**
     * 录制并返回WAV格式的完整数据（包含WAV头）
     * @param durationSeconds 录音时长（秒）
     * @return 包含WAV头的完整音频数据
     * @throws Exception 录音异常
     */
    public byte[] recordAsWavBytes(int durationSeconds) throws Exception {
        byte[] rawAudioData = record(durationSeconds);
        return addWavHeader(rawAudioData);
    }

    /**
     * 为原始音频数据添加WAV文件头
     * @param audioData 原始音频数据
     * @return 包含WAV头的完整音频数据
     */
    private byte[] addWavHeader(byte[] audioData) {
        if (format == null) {
            throw new IllegalStateException("音频格式未设置");
        }

        long totalDataLen = audioData.length + 36; // 数据长度 + 头文件长度
        long sampleRate = (long) format.getSampleRate();
        int channels = format.getChannels();
        int bitsPerSample = format.getSampleSizeInBits();
        long byteRate = sampleRate * channels * bitsPerSample / 8;
        int blockAlign = channels * bitsPerSample / 8;

        ByteArrayOutputStream header = new ByteArrayOutputStream();

        try {
            // RIFF头
            header.write("RIFF".getBytes()); // ChunkID
            header.write(intToByteArray((int) totalDataLen), 0, 4); // ChunkSize
            header.write("WAVE".getBytes()); // Format

            // fmt子块
            header.write("fmt ".getBytes()); // Subchunk1ID
            header.write(intToByteArray(16), 0, 4); // Subchunk1Size (16 for PCM)
            header.write(shortToByteArray((short) 1), 0, 2); // AudioFormat (1 for PCM)
            header.write(shortToByteArray((short) channels), 0, 2); // NumChannels
            header.write(intToByteArray((int) sampleRate), 0, 4); // SampleRate
            header.write(intToByteArray((int) byteRate), 0, 4); // ByteRate
            header.write(shortToByteArray((short) blockAlign), 0, 2); // BlockAlign
            header.write(shortToByteArray((short) bitsPerSample), 0, 2); // BitsPerSample

            // data子块
            header.write("data".getBytes()); // Subchunk2ID
            header.write(intToByteArray(audioData.length), 0, 4); // Subchunk2Size

            // 合并头文件和音频数据
            ByteArrayOutputStream wavFile = new ByteArrayOutputStream();
            wavFile.write(header.toByteArray());
            wavFile.write(audioData);

            return wavFile.toByteArray();
        } catch (IOException e) {
            throw new RuntimeException("生成WAV头时出错", e);
        }
    }

    /**
     * int转换为4字节小端序
     */
    private byte[] intToByteArray(int data) {
        return new byte[] {
                (byte) (data & 0xff),
                (byte) ((data >> 8) & 0xff),
                (byte) ((data >> 16) & 0xff),
                (byte) ((data >> 24) & 0xff)
        };
    }

    /**
     * short转换为2字节小端序
     */
    private byte[] shortToByteArray(short data) {
        return new byte[] {
                (byte) (data & 0xff),
                (byte) ((data >> 8) & 0xff)
        };
    }

    public void stopRecording() {
        isRecording = false;
        if (line != null) {
            line.stop();
            line.close();
        }
    }
}